import logging
from logging.handlers import RotatingFileHandler
import os
from datetime import datetime

def setup_logging(base_log_dir='logs', log_prefix='train_'):
    """
    统一的日志系统设置
    返回: (logger, log_dir) 元组
    """
    # 创建按时间命名的日志目录
    current_time = datetime.now().strftime('%Y%m%d_%H%M%S')
    log_dir = os.path.join(base_log_dir, current_time)
    
    # 创建所有必要的子目录
    os.makedirs(os.path.join(log_dir, 'models'), exist_ok=True)
    os.makedirs(os.path.join(log_dir, 'plots'), exist_ok=True)
    
    # 配置日志文件
    log_filename = f"{log_prefix}{current_time}.log"
    log_filepath = os.path.join(log_dir, log_filename)
    
    # 初始化日志记录器
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)
    
    # 清除现有处理器（避免重复日志）
    for handler in logger.handlers[:]:
        logger.removeHandler(handler)
    
    # 文件处理器（按大小轮转）
    file_handler = RotatingFileHandler(
        log_filepath, 
        maxBytes=10 * 1024 * 1024,  # 10MB
        backupCount=5
    )
    file_handler.setFormatter(logging.Formatter(
        '%(asctime)s - %(levelname)s - %(message)s'
    ))
    
    # 控制台处理器
    console_handler = logging.StreamHandler()
    console_handler.setFormatter(logging.Formatter(
        '%(asctime)s - %(levelname)s - %(message)s'
    ))
    
    logger.addHandler(file_handler)
    logger.addHandler(console_handler)
    
    return logger, log_dir